package main

import (
	"encoding/binary"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"sort"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

// 新增：Prometheus指标定义
var (
	// 流量统计指标
	flowBytesSent = promauto.NewCounterVec(
		prometheus.CounterOpts{
			Name: "netflow_bytes_sent_total",
			Help: "Total bytes sent in network flows",
		},
		[]string{"src_ip", "dst_ip", "src_port", "dst_port", "protocol", "app_name"},
	)

	flowBytesReceived = promauto.NewCounterVec(
		prometheus.CounterOpts{
			Name: "netflow_bytes_received_total",
			Help: "Total bytes received in network flows",
		},
		[]string{"src_ip", "dst_ip", "src_port", "dst_port", "protocol", "app_name"},
	)

	flowPacketsSent = promauto.NewCounterVec(
		prometheus.CounterOpts{
			Name: "netflow_packets_sent_total",
			Help: "Total packets sent in network flows",
		},
		[]string{"src_ip", "dst_ip", "src_port", "dst_port", "protocol", "app_name"},
	)

	flowPacketsReceived = promauto.NewCounterVec(
		prometheus.CounterOpts{
			Name: "netflow_packets_received_total",
			Help: "Total packets received in network flows",
		},
		[]string{"src_ip", "dst_ip", "src_port", "dst_port", "protocol", "app_name"},
	)

	// TCP特定指标
	flowTcpRetransmissions = promauto.NewCounterVec(
		prometheus.CounterOpts{
			Name: "netflow_tcp_retransmissions_total",
			Help: "Total TCP retransmissions",
		},
		[]string{"src_ip", "dst_ip", "src_port", "dst_port", "app_name"},
	)

	flowTcpLostPackets = promauto.NewCounterVec(
		prometheus.CounterOpts{
			Name: "netflow_tcp_lost_packets_total",
			Help: "Total TCP lost packets",
		},
		[]string{"src_ip", "dst_ip", "src_port", "dst_port", "app_name"},
	)

	flowTcpRtt = promauto.NewGaugeVec(
		prometheus.GaugeOpts{
			Name: "netflow_tcp_rtt_seconds",
			Help: "Average TCP round-trip time in seconds",
		},
		[]string{"src_ip", "dst_ip", "src_port", "dst_port", "app_name"},
	)

	// 流状态指标
	flowDuration = promauto.NewGaugeVec(
		prometheus.GaugeOpts{
			Name: "netflow_duration_seconds",
			Help: "Duration of network flows in seconds",
		},
		[]string{"src_ip", "dst_ip", "src_port", "dst_port", "protocol", "status", "app_name"},
	)

	activeFlows = promauto.NewGaugeVec(
		prometheus.GaugeOpts{
			Name: "netflow_active_flows",
			Help: "Number of active network flows",
		},
		[]string{"protocol", "app_category"},
	)
)

// NetFlow 统一网络流结构，支持TCP/UDP/ICMP
type NetFlow struct {
	// 基础标识（兼容IPv4/IPv6）
	SrcIP   string
	SrcPort uint16 // UDP/TCP用，ICMP填0
	DstIP   string
	DstPort uint16 // UDP/TCP用，ICMP填0
	Proto   string // TCP/UDP/ICMP

	// 连接/会话状态
	StartTime time.Time
	LastSeen  time.Time
	Status    string // TCP: ESTABLISHED/SYN_SENT等；UDP/ICMP: ACTIVE/INACTIVE

	// 流量统计（通用）
	BytesSent       uint64
	BytesReceived   uint64
	PacketsSent     uint64
	PacketsReceived uint64
	AvgPacketSize   float64

	// TCP专属指标（UDP/ICMP为空）
	RetransmittedPackets uint64
	DuplicateAcks        uint64
	LostPackets          uint64
	MinRTT               time.Duration
	MaxRTT               time.Duration
	AvgRTT               time.Duration
	RTTSamples           []time.Duration
	LastSentSeq          uint32
	LastAckSeq           uint32
	LastReceivedSeq      uint32
	LastSentTime         time.Time

	// UDP专属指标
	UDPDatagramLoss uint64 // 基于序列号推断（需应用层支持）

	// ICMP专属指标
	ICMPType     uint8  // ICMP类型（如8=请求，0=响应）
	ICMPCode     uint8  // ICMP代码
	ICMPSequence uint16 // ICMP序列号（部分类型支持）

	// 应用识别相关（通用）
	AppName     string  // 应用名称
	AppCategory string  // 应用类别
	Confidence  float64 // 识别可信度(0-1)

	// Payload存储（通用）
	PayloadsSent     [][]byte // 源到目标的payload
	PayloadsReceived [][]byte // 目标到源的payload
	MaxPayloadsCount int      // 最大存储的payload数量
	MaxPayloadSize   int      // 单个payload的最大存储字节数

	// 新增：用于Prometheus指标增量更新的计数器
	lastBytesSent       uint64
	lastBytesReceived   uint64
	lastPacketsSent     uint64
	lastPacketsReceived uint64
	lastRetransmissions uint64
	lastLostPackets     uint64
}

// Duration 计算连接持续时长
func (f *NetFlow) Duration() time.Duration {
	return f.LastSeen.Sub(f.StartTime)
}

// FlowKey 统一流标识，兼容不同协议
type FlowKey struct {
	SrcIP   string
	SrcPort uint16
	DstIP   string
	DstPort uint16
	Proto   string // 区分TCP/UDP/ICMP
}

// 全局配置和状态
var (
	netFlows          = make(map[FlowKey]*NetFlow)
	flowMutex         = &sync.RWMutex{}
	iface             string
	snapshotLen       int32         = 1600
	promiscuous       bool          = false
	timeout           time.Duration = pcap.BlockForever
	printInterval     time.Duration = 5 * time.Second
	maxPayloadsCount  int           = 100              // 每个方向默认最多存100个payload
	maxPayloadSize    int           = 4096             // 单个payload默认最多存4096字节
	payloadStorageDir string        = "net_payloads"   // payload存储目录
	udpTimeout        time.Duration = 30 * time.Second // UDP会话超时时间
	icmpTimeout       time.Duration = 10 * time.Second // ICMP会话超时时间
	displayCountMax   int           = 20
	metricsAddr       string        = ":9091" // 新增：Prometheus指标暴露地址
)

func init() {
	flag.StringVar(&iface, "iface", "", "网络接口名称 (例如: eth0, en0)")
	flag.DurationVar(&printInterval, "interval", 5*time.Second, "统计信息打印间隔")
	flag.IntVar(&displayCountMax, "display", 20, "统计信息打印行数")
	flag.IntVar(&maxPayloadsCount, "max-payloads", 100, "每个方向最大存储的payload数量")
	flag.IntVar(&maxPayloadSize, "max-payload-size", 4096, "单个payload最大存储字节数")
	flag.StringVar(&payloadStorageDir, "payload-dir", "net_payloads", "payload文件存储目录")
	flag.StringVar(&metricsAddr, "metrics-addr", ":9091", "Prometheus指标暴露地址") // 新增
	flag.Parse()

	if iface == "" {
		log.Fatal("必须指定网络接口 (-iface)")
	}

	// 创建payload存储目录
	if err := os.MkdirAll(payloadStorageDir, 0755); err != nil {
		log.Fatalf("无法创建payload存储目录: %v", err)
	}

	// 新增：启动Prometheus指标HTTP服务
	go func() {
		http.Handle("/metrics", promhttp.Handler())
		log.Printf("Prometheus指标暴露在: http://%s/metrics", metricsAddr)
		if err := http.ListenAndServe(metricsAddr, nil); err != nil {
			log.Fatalf("启动指标服务失败: %v", err)
		}
	}()
}

// getFlowKey 生成统一流标识，确保双向流唯一
func getFlowKey(srcIP net.IP, srcPort uint16, dstIP net.IP, dstPort uint16, proto string) FlowKey {
	srcStr := fmt.Sprintf("%s:%d", srcIP, srcPort)
	dstStr := fmt.Sprintf("%s:%d", dstIP, dstPort)

	// 按字符串排序确保双向流key一致
	if srcStr < dstStr {
		return FlowKey{
			SrcIP:   srcIP.String(),
			SrcPort: srcPort,
			DstIP:   dstIP.String(),
			DstPort: dstPort,
			Proto:   proto,
		}
	}
	return FlowKey{
		SrcIP:   dstIP.String(),
		SrcPort: dstPort,
		DstIP:   srcIP.String(),
		DstPort: srcPort,
		Proto:   proto,
	}
}

// 新增：更新Prometheus指标
func updatePrometheusMetrics(flow *NetFlow) {
	// 通用标签
	labels := prometheus.Labels{
		"src_ip":   flow.SrcIP,
		"dst_ip":   flow.DstIP,
		"src_port": fmt.Sprintf("%d", flow.SrcPort),
		"dst_port": fmt.Sprintf("%d", flow.DstPort),
		"protocol": flow.Proto,
		"app_name": flow.AppName,
	}

	// 流量增量更新
	if flow.BytesSent > flow.lastBytesSent {
		flowBytesSent.With(labels).Add(float64(flow.BytesSent - flow.lastBytesSent))
		flow.lastBytesSent = flow.BytesSent
	}

	if flow.BytesReceived > flow.lastBytesReceived {
		flowBytesReceived.With(labels).Add(float64(flow.BytesReceived - flow.lastBytesReceived))
		flow.lastBytesReceived = flow.BytesReceived
	}

	if flow.PacketsSent > flow.lastPacketsSent {
		flowPacketsSent.With(labels).Add(float64(flow.PacketsSent - flow.lastPacketsSent))
		flow.lastPacketsSent = flow.PacketsSent
	}

	if flow.PacketsReceived > flow.lastPacketsReceived {
		flowPacketsReceived.With(labels).Add(float64(flow.PacketsReceived - flow.lastPacketsReceived))
		flow.lastPacketsReceived = flow.PacketsReceived
	}

	// TCP特定指标
	if flow.Proto == "TCP" {
		tcpLabels := prometheus.Labels{
			"src_ip":   flow.SrcIP,
			"dst_ip":   flow.DstIP,
			"src_port": fmt.Sprintf("%d", flow.SrcPort),
			"dst_port": fmt.Sprintf("%d", flow.DstPort),
			"app_name": flow.AppName,
		}

		if flow.RetransmittedPackets > flow.lastRetransmissions {
			flowTcpRetransmissions.With(tcpLabels).Add(float64(flow.RetransmittedPackets - flow.lastRetransmissions))
			flow.lastRetransmissions = flow.RetransmittedPackets
		}

		if flow.LostPackets > flow.lastLostPackets {
			flowTcpLostPackets.With(tcpLabels).Add(float64(flow.LostPackets - flow.lastLostPackets))
			flow.lastLostPackets = flow.LostPackets
		}

		if flow.AvgRTT > 0 {
			flowTcpRtt.With(tcpLabels).Set(flow.AvgRTT.Seconds())
		}
	}

	// 流持续时间
	durationLabels := prometheus.Labels{
		"src_ip":   flow.SrcIP,
		"dst_ip":   flow.DstIP,
		"src_port": fmt.Sprintf("%d", flow.SrcPort),
		"dst_port": fmt.Sprintf("%d", flow.DstPort),
		"protocol": flow.Proto,
		"status":   flow.Status,
		"app_name": flow.AppName,
	}
	flowDuration.With(durationLabels).Set(flow.Duration().Seconds())
}

// 新增：定期更新活跃流指标
func updateActiveFlowsMetrics() {
	flowMutex.RLock()
	defer flowMutex.RUnlock()

	// 先重置活跃流指标
	activeFlows.Reset()

	// 统计各类型活跃流数量
	counter := make(map[string]map[string]int) // protocol -> app_category -> count
	for _, flow := range netFlows {
		if _, ok := counter[flow.Proto]; !ok {
			counter[flow.Proto] = make(map[string]int)
		}
		counter[flow.Proto][flow.AppCategory]++
	}

	// 更新指标
	for proto, cats := range counter {
		for cat, count := range cats {
			activeFlows.WithLabelValues(proto, cat).Set(float64(count))
		}
	}
}

// savePayload 保存payload到内存（控制数量和大小）
func savePayload(payload []byte, payloads *[][]byte, maxCount, maxSize int) {
	if len(payload) == 0 {
		return
	}

	// 限制单个payload大小
	if len(payload) > maxSize {
		payload = payload[:maxSize]
	}

	// 限制总数量（FIFO）
	if len(*payloads) >= maxCount {
		*payloads = (*payloads)[1:]
	}

	// 保存副本避免引用污染
	payloadCopy := make([]byte, len(payload))
	copy(payloadCopy, payload)
	*payloads = append(*payloads, payloadCopy)
}

// savePayloadsToFile 将内存中的payload保存到文件
func savePayloadsToFile(flow *NetFlow) error {
	if len(flow.PayloadsSent) == 0 && len(flow.PayloadsReceived) == 0 {
		return nil
	}

	// 创建流专属目录（包含协议标识）
	flowDir := fmt.Sprintf("%s/%s_%s_%d-%s_%d",
		payloadStorageDir,
		flow.Proto,
		flow.SrcIP, flow.SrcPort,
		flow.DstIP, flow.DstPort)
	if err := os.MkdirAll(flowDir, 0755); err != nil {
		return err
	}

	// 保存发送的payload
	for i, payload := range flow.PayloadsSent {
		filename := fmt.Sprintf("%s/sent_%d_%d.bin",
			flowDir,
			flow.StartTime.UnixNano(),
			i)
		if err := os.WriteFile(filename, payload, 0644); err != nil {
			log.Printf("保存发送的payload失败: %v", err)
		}
	}

	// 保存接收的payload
	for i, payload := range flow.PayloadsReceived {
		filename := fmt.Sprintf("%s/received_%d_%d.bin",
			flowDir,
			flow.StartTime.UnixNano(),
			i)
		if err := os.WriteFile(filename, payload, 0644); err != nil {
			log.Printf("保存接收的payload失败: %v", err)
		}
	}

	return nil
}

// updateTCPRTT 更新TCP流的RTT信息
func updateTCPRTT(flow *NetFlow, tcp *layers.TCP, packetTime time.Time) {
	if tcp.ACK && tcp.Ack > 0 && flow.LastSentSeq > 0 && tcp.Ack >= flow.LastSentSeq {
		rtt := packetTime.Sub(flow.LastSentTime)
		if rtt > 0 {
			// 限制RTT样本数量
			if len(flow.RTTSamples) > 100 {
				flow.RTTSamples = flow.RTTSamples[1:]
			}
			flow.RTTSamples = append(flow.RTTSamples, rtt)

			// 更新RTT统计
			if flow.MinRTT == 0 || rtt < flow.MinRTT {
				flow.MinRTT = rtt
			}
			if rtt > flow.MaxRTT {
				flow.MaxRTT = rtt
			}

			total := time.Duration(0)
			for _, t := range flow.RTTSamples {
				total += t
			}
			flow.AvgRTT = total / time.Duration(len(flow.RTTSamples))
		}
	}
}

// extractTLSSNI 从TLS握手数据中提取SNI
func extractTLSSNI(payload []byte) (string, bool) {
	if len(payload) < 5 || payload[0] != 0x16 {
		return "", false
	}

	pos := 5
	if pos+1 > len(payload) {
		return "", false
	}

	handshakeType := payload[pos]
	pos++
	if handshakeType != 0x01 {
		return "", false
	}

	if pos+3 > len(payload) {
		return "", false
	}
	handshakeLen := int(binary.BigEndian.Uint32(append([]byte{0}, payload[pos:]...)))
	pos += 3
	if pos+handshakeLen > len(payload) {
		return "", false
	}

	pos += 2 + 32 // 跳过版本号(2)和随机数(32)
	if pos+1 > len(payload) {
		return "", false
	}

	sessionIDLen := int(payload[pos])
	pos++
	pos += sessionIDLen
	if pos+2 > len(payload) {
		return "", false
	}

	cipherSuitesLen := int(binary.BigEndian.Uint16(payload[pos:]))
	pos += 2
	pos += cipherSuitesLen
	if pos+1 > len(payload) {
		return "", false
	}

	compressionMethodsLen := int(payload[pos])
	pos++
	pos += compressionMethodsLen
	if pos+2 > len(payload) {
		return "", false
	}

	extensionsLen := int(binary.BigEndian.Uint16(payload[pos:]))
	pos += 2
	end := pos + extensionsLen

	for pos < end {
		if pos+4 > len(payload) {
			break
		}
		extType := binary.BigEndian.Uint16(payload[pos:])
		extLen := int(binary.BigEndian.Uint16(payload[pos+2:]))
		pos += 4

		if extType == 0 { // SNI扩展类型
			if pos+2 > len(payload) {
				break
			}
			sniListLen := int(binary.BigEndian.Uint16(payload[pos:]))
			pos += 2
			sniEnd := pos + sniListLen

			for pos < sniEnd {
				if pos+3 > len(payload) {
					break
				}
				nameType := payload[pos]
				nameLen := int(binary.BigEndian.Uint16(payload[pos+1:]))
				pos += 3

				if nameType == 0 && pos+nameLen <= len(payload) {
					return string(payload[pos : pos+nameLen]), true
				}
				pos += nameLen
			}
			break
		}
		pos += extLen
	}

	return "", false
}

// extractAppFeatures 从应用层数据中提取特征
func extractAppFeatures(payload []byte) map[string]string {
	features := make(map[string]string)
	if len(payload) == 0 {
		return features
	}

	firstLine := string(payload[:min(len(payload), 512)])
	if strings.HasPrefix(firstLine, "GET ") || strings.HasPrefix(firstLine, "POST ") ||
		strings.HasPrefix(firstLine, "PUT ") || strings.HasPrefix(firstLine, "DELETE ") ||
		strings.HasPrefix(firstLine, "HEAD ") || strings.HasPrefix(firstLine, "OPTIONS ") ||
		strings.HasPrefix(firstLine, "HTTP/1.") {

		if hostStart := strings.Index(firstLine, "Host: "); hostStart > 0 {
			hostEnd := strings.Index(firstLine[hostStart:], "\r\n")
			if hostEnd > 0 {
				host := strings.TrimSpace(firstLine[hostStart+6 : hostStart+hostEnd])
				features["http_host"] = host
			}
		}

		if uaStart := strings.Index(firstLine, "User-Agent: "); uaStart > 0 {
			uaEnd := strings.Index(firstLine[uaStart:], "\r\n")
			if uaEnd > 0 {
				ua := strings.TrimSpace(firstLine[uaStart+12 : uaStart+uaEnd])
				features["http_ua"] = ua
			}
		}
	}

	if len(payload) > 5 && payload[0] == 0x16 {
		if sni, ok := extractTLSSNI(payload); ok {
			features["tls_sni"] = sni
		}
	}

	return features
}

// identifyByPortProtocol 基于端口和协议的基础识别
func identifyByPortProtocol(flow *NetFlow) (appName string, confidence float64) {
	// 按协议区分端口映射
	portAppMap := map[string]map[uint16]string{
		"TCP": {
			80:   "HTTP服务",
			443:  "HTTPS服务",
			21:   "FTP文件传输服务",
			22:   "SSH远程登录服务",
			23:   "Telnet远程登录服务",
			25:   "SMTP邮件发送服务",
			53:   "DNS域名解析服务（TCP模式）",
			8080: "HTTP代理/备用Web服务",
			8443: "HTTPS备用服务",
			110:  "POP3邮件接收服务",
			143:  "IMAP邮件接收服务",
			3306: "MySQL数据库服务",
			3389: "Windows远程桌面（RDP）",
			5432: "PostgreSQL数据库服务",
			5900: "VNC远程控制服务",
			6379: "Redis缓存服务",
		},
		"UDP": {
			53:   "DNS服务",
			67:   "DHCP服务",
			68:   "DHCP客户端",
			123:  "NTP服务",
			5060: "SIP服务",
			3478: "ICE服务",
		},
		"ICMP": {
			0: "ICMP响应",
			8: "ICMP请求",
		},
	}

	// TCP/UDP用端口匹配
	if flow.Proto == "TCP" || flow.Proto == "UDP" {
		if appMap, ok := portAppMap[flow.Proto]; ok {
			if app, ok := appMap[flow.SrcPort]; ok {
				return app, 0.8
			}
			if app, ok := appMap[flow.DstPort]; ok {
				return app, 0.8
			}
		}
	}

	// ICMP用类型匹配
	if flow.Proto == "ICMP" {
		if appMap, ok := portAppMap["ICMP"]; ok {
			if app, ok := appMap[uint16(flow.ICMPType)]; ok {
				return app, 0.9
			}
		}
	}

	return "未知服务", 0.2
}

// identifyByDomainFeatures 基于域名和特征字符串识别应用
func identifyByDomainFeatures(features map[string]string) (appName, appCategory string, confidence float64) {
	domainMap := map[string][2]string{
		"baidu.com":     {"百度搜索", "网页浏览"},
		"google.com":    {"谷歌搜索", "网页浏览"},
		"weixin.qq.com": {"微信", "即时通讯"},
		"qq.com":        {"腾讯QQ", "即时通讯"},
		"youtube.com":   {"YouTube", "视频流"},
		"bilibili.com":  {"哔哩哔哩", "视频流"},
		"taobao.com":    {"淘宝", "电子商务"},
		"jd.com":        {"京东", "电子商务"},
		"aliyun.com":    {"阿里云", "云服务"},
		"github.com":    {"GitHub", "代码托管"},
	}

	var targetDomain string
	if host, ok := features["http_host"]; ok {
		targetDomain = host
	} else if sni, ok := features["tls_sni"]; ok {
		targetDomain = sni
	}

	if targetDomain != "" {
		for domain, appInfo := range domainMap {
			if strings.Contains(targetDomain, domain) {
				return appInfo[0], appInfo[1], 0.9
			}
		}
	}

	if ua, ok := features["http_ua"]; ok {
		switch {
		case strings.Contains(ua, "Chrome"):
			return "Chrome浏览器", "网页浏览", 0.7
		case strings.Contains(ua, "Firefox"):
			return "Firefox浏览器", "网页浏览", 0.7
		case strings.Contains(ua, "Safari"):
			return "Safari浏览器", "网页浏览", 0.7
		case strings.Contains(ua, "Edge"):
			return "Edge浏览器", "网页浏览", 0.7
		case strings.Contains(ua, "Mobile"):
			return "移动设备浏览器", "网页浏览", 0.6
		}
	}

	return "未知应用", "未知类别", 0.3
}

// identifyByBehavior 基于流量行为特征识别应用类别
func identifyByBehavior(flow *NetFlow) (appCategory string, confidence float64) {
	duration := flow.Duration().Seconds()
	if duration < 1 {
		return "未知类别", 0.2
	}

	totalBytes := flow.BytesSent + flow.BytesReceived
	bytesPerSec := float64(totalBytes) / duration
	packetsPerSec := float64(flow.PacketsSent+flow.PacketsReceived) / duration

	// 针对不同协议调整行为识别规则
	switch flow.Proto {
	case "TCP":
		switch {
		case bytesPerSec > 1024*1024 && packetsPerSec > 10:
			return "视频流", 0.7
		case bytesPerSec < 1024*100 && packetsPerSec > 5 && flow.AvgRTT < 500*time.Millisecond:
			return "即时通讯", 0.6
		}
	case "UDP":
		switch {
		case packetsPerSec > 50 && flow.AvgPacketSize < 1000:
			return "实时通讯", 0.6 // 如语音通话
		case bytesPerSec > 512*1024 && flow.AvgPacketSize > 1400:
			return "视频流", 0.7 // 如UDP视频传输
		case flow.DstPort == 53 || flow.SrcPort == 53:
			return "域名解析", 0.9
		}
	case "ICMP":
		return "网络探测", 0.8 // ICMP多为ping等探测行为
	}

	if bytesPerSec < 1024*500 && packetsPerSec < 100 {
		return "网页浏览", 0.5
	}

	return "未知类别", 0.2
}

// identifyApplication 综合识别应用
func identifyApplication(flow *NetFlow, payload []byte) {
	features := extractAppFeatures(payload)

	portApp, portConf := identifyByPortProtocol(flow)
	domainApp, domainCat, domainConf := identifyByDomainFeatures(features)
	behaviorCat, behaviorConf := identifyByBehavior(flow)

	// 按协议调整权重（ICMP优先行为识别，TCP/UDP优先域名/端口）
	switch flow.Proto {
	case "ICMP":
		flow.AppName = portApp
		flow.AppCategory = behaviorCat
		flow.Confidence = (portConf + behaviorConf) / 2
	case "UDP":
		if domainConf > portConf {
			flow.AppName = domainApp
			flow.AppCategory = domainCat
			flow.Confidence = domainConf
		} else {
			flow.AppName = portApp
			flow.AppCategory = behaviorCat
			flow.Confidence = (portConf + behaviorConf) / 2
		}
	default: // TCP
		if domainConf >= portConf && domainConf >= behaviorConf {
			flow.AppName = domainApp
			flow.AppCategory = domainCat
			flow.Confidence = domainConf
		} else if portConf >= domainConf && portConf >= behaviorConf {
			flow.AppName = portApp
			flow.AppCategory = "基础服务"
			flow.Confidence = portConf
		} else {
			flow.AppName = "未知应用"
			flow.AppCategory = behaviorCat
			flow.Confidence = behaviorConf
		}
	}
}

// handleTCPPacket 处理TCP数据包
func handleTCPPacket(srcIp net.IP, dstIp net.IP, tcp *layers.TCP, packet gopacket.Packet) {
	srcPort, dstPort := uint16(tcp.SrcPort), uint16(tcp.DstPort)
	key := getFlowKey(srcIp, srcPort, dstIp, dstPort, "TCP")
	isOriginalDirection := (key.SrcIP == srcIp.String() && key.SrcPort == srcPort)

	flowMutex.Lock()
	defer flowMutex.Unlock()

	flow, exists := netFlows[key]
	if !exists {
		flow = &NetFlow{
			SrcIP:            key.SrcIP,
			SrcPort:          key.SrcPort,
			DstIP:            key.DstIP,
			DstPort:          key.DstPort,
			Proto:            "TCP",
			StartTime:        packet.Metadata().Timestamp,
			Status:           "NEW",
			MaxPayloadsCount: maxPayloadsCount,
			MaxPayloadSize:   maxPayloadSize,
			PayloadsSent:     make([][]byte, 0),
			PayloadsReceived: make([][]byte, 0),
		}
		netFlows[key] = flow
	}

	flow.LastSeen = packet.Metadata().Timestamp

	// 更新TCP状态
	if tcp.SYN && !tcp.ACK {
		flow.Status = "SYN_SENT"
	} else if tcp.SYN && tcp.ACK {
		flow.Status = "SYN_RECEIVED"
	} else if tcp.ACK && flow.Status != "ESTABLISHED" {
		flow.Status = "ESTABLISHED"
	} else if tcp.FIN {
		flow.Status = "FIN_WAIT"
	} else if tcp.RST {
		flow.Status = "RESET"
	}

	// 更新流量统计和payload
	packetSize := uint64(len(packet.Data()))
	payload := tcp.Payload
	if isOriginalDirection {
		flow.BytesSent += packetSize
		flow.PacketsSent++
		flow.LastSentSeq = tcp.Seq
		flow.LastSentTime = packet.Metadata().Timestamp
		if len(payload) > 0 {
			savePayload(payload, &flow.PayloadsSent, flow.MaxPayloadsCount, flow.MaxPayloadSize)
		}
	} else {
		flow.BytesReceived += packetSize
		flow.PacketsReceived++
		flow.LastReceivedSeq = tcp.Seq
		if len(payload) > 0 {
			savePayload(payload, &flow.PayloadsReceived, flow.MaxPayloadsCount, flow.MaxPayloadSize)
		}
	}

	// 计算平均包大小
	totalBytes := flow.BytesSent + flow.BytesReceived
	totalPackets := flow.PacketsSent + flow.PacketsReceived
	if totalPackets > 0 {
		flow.AvgPacketSize = float64(totalBytes) / float64(totalPackets)
	}

	if tcp.ACK {
		flow.LastAckSeq = tcp.Ack
	}

	// 检测重传和丢包
	if !tcp.SYN && isOriginalDirection {
		if tcp.Seq <= flow.LastAckSeq && flow.PacketsSent > 1 {
			flow.RetransmittedPackets++
		}
	} else if !tcp.SYN && !isOriginalDirection {
		if tcp.ACK && tcp.Ack == flow.LastAckSeq && flow.PacketsReceived > 1 {
			flow.DuplicateAcks++
			flow.LostPackets++
		}
	}

	// 更新RTT
	updateTCPRTT(flow, tcp, packet.Metadata().Timestamp)

	// 应用识别
	if len(payload) > 0 {
		identifyApplication(flow, payload)
	}

	// 新增：更新Prometheus指标
	updatePrometheusMetrics(flow)
}

// handleUDPPacket 处理UDP数据包
func handleUDPPacket(srcIp net.IP, dstIp net.IP, udp *layers.UDP, packet gopacket.Packet) {
	srcPort, dstPort := uint16(udp.SrcPort), uint16(udp.DstPort)
	key := getFlowKey(srcIp, srcPort, dstIp, dstPort, "UDP")
	isOriginalDirection := (key.SrcIP == srcIp.String() && key.SrcPort == srcPort)

	flowMutex.Lock()
	defer flowMutex.Unlock()

	flow, exists := netFlows[key]
	if !exists {
		flow = &NetFlow{
			SrcIP:            key.SrcIP,
			SrcPort:          key.SrcPort,
			DstIP:            key.DstIP,
			DstPort:          key.DstPort,
			Proto:            "UDP",
			StartTime:        packet.Metadata().Timestamp,
			Status:           "ACTIVE",
			MaxPayloadsCount: maxPayloadsCount,
			MaxPayloadSize:   maxPayloadSize,
			PayloadsSent:     make([][]byte, 0),
			PayloadsReceived: make([][]byte, 0),
		}
		netFlows[key] = flow
	}

	flow.LastSeen = packet.Metadata().Timestamp

	// 更新UDP状态（活跃/非活跃）
	if time.Since(flow.LastSeen) < udpTimeout/2 {
		flow.Status = "ACTIVE"
	} else {
		flow.Status = "INACTIVE"
	}

	// 更新流量统计和payload
	packetSize := uint64(len(packet.Data()))
	payload := udp.Payload
	if isOriginalDirection {
		flow.BytesSent += packetSize
		flow.PacketsSent++
		if len(payload) > 0 {
			savePayload(payload, &flow.PayloadsSent, flow.MaxPayloadsCount, flow.MaxPayloadSize)
		}
	} else {
		flow.BytesReceived += packetSize
		flow.PacketsReceived++
		if len(payload) > 0 {
			savePayload(payload, &flow.PayloadsReceived, flow.MaxPayloadsCount, flow.MaxPayloadSize)
		}
	}

	// 计算平均包大小
	totalBytes := flow.BytesSent + flow.BytesReceived
	totalPackets := flow.PacketsSent + flow.PacketsReceived
	if totalPackets > 0 {
		flow.AvgPacketSize = float64(totalBytes) / float64(totalPackets)
	}

	// 应用识别
	if len(payload) > 0 {
		identifyApplication(flow, payload)
	}

	// 新增：更新Prometheus指标
	updatePrometheusMetrics(flow)
}

// handleICMPPacket 处理ICMP数据包（IPv4）
func handleICMPPacket(ip *layers.IPv4, icmp *layers.ICMPv4, packet gopacket.Packet) {
	key := getFlowKey(ip.SrcIP, 0, ip.DstIP, 0, "ICMP") // ICMP无端口，用0填充
	isOriginalDirection := (key.SrcIP == ip.SrcIP.String())

	flowMutex.Lock()
	defer flowMutex.Unlock()

	flow, exists := netFlows[key]
	if !exists {
		flow = &NetFlow{
			SrcIP:            key.SrcIP,
			SrcPort:          0,
			DstIP:            key.DstIP,
			DstPort:          0,
			Proto:            "ICMP",
			StartTime:        packet.Metadata().Timestamp,
			Status:           "ACTIVE",
			ICMPType:         icmp.TypeCode.Type(),
			ICMPCode:         icmp.TypeCode.Code(),
			MaxPayloadsCount: maxPayloadsCount,
			MaxPayloadSize:   maxPayloadSize,
			PayloadsSent:     make([][]byte, 0),
			PayloadsReceived: make([][]byte, 0),
		}
		// 提取ICMP序列号（仅部分类型有）
		if icmp.TypeCode.Type() == layers.ICMPv4TypeEchoRequest ||
			icmp.TypeCode.Type() == layers.ICMPv4TypeEchoReply {
			flow.ICMPSequence = binary.BigEndian.Uint16(icmp.Payload[:2])
		}
		netFlows[key] = flow
	}

	flow.LastSeen = packet.Metadata().Timestamp

	// 更新ICMP状态
	if time.Since(flow.LastSeen) < icmpTimeout/2 {
		flow.Status = "ACTIVE"
	} else {
		flow.Status = "INACTIVE"
	}

	// 更新流量统计和payload
	packetSize := uint64(len(packet.Data()))
	payload := icmp.Payload
	if isOriginalDirection {
		flow.BytesSent += packetSize
		flow.PacketsSent++
		if len(payload) > 0 {
			savePayload(payload, &flow.PayloadsSent, flow.MaxPayloadsCount, flow.MaxPayloadSize)
		}
	} else {
		flow.BytesReceived += packetSize
		flow.PacketsReceived++
		if len(payload) > 0 {
			savePayload(payload, &flow.PayloadsReceived, flow.MaxPayloadsCount, flow.MaxPayloadSize)
		}
	}

	// 计算平均包大小
	totalBytes := flow.BytesSent + flow.BytesReceived
	totalPackets := flow.PacketsSent + flow.PacketsReceived
	if totalPackets > 0 {
		flow.AvgPacketSize = float64(totalBytes) / float64(totalPackets)
	}

	// 应用识别（ICMP类型固定）
	identifyApplication(flow, payload)

	// 新增：更新Prometheus指标
	updatePrometheusMetrics(flow)
}

// 处理IPv6数据包
func handleIPv6Packet(ip6 *layers.IPv6, packet gopacket.Packet) {
	switch ip6.NextHeader {
	case layers.IPProtocolTCP:
		tcp := &layers.TCP{}
		if err := tcp.DecodeFromBytes(ip6.Payload, gopacket.NilDecodeFeedback); err == nil {
			handleTCPPacket(ip6.SrcIP, ip6.DstIP, tcp, packet)
		}
	case layers.IPProtocolUDP:
		udp := &layers.UDP{}
		if err := udp.DecodeFromBytes(ip6.Payload, gopacket.NilDecodeFeedback); err == nil {
			handleUDPPacket(ip6.SrcIP, ip6.DstIP, udp, packet)
		}
	case layers.IPProtocolICMPv6:
		icmp6 := &layers.ICMPv6{}
		if err := icmp6.DecodeFromBytes(ip6.Payload, gopacket.NilDecodeFeedback); err == nil {
			key := getFlowKey(ip6.SrcIP, 0, ip6.DstIP, 0, "ICMPv6")
			fmt.Println(key)
		}
	}
}

// 处理所有数据包
func handlePacket(packet gopacket.Packet) {
	// 解析IPv4
	if ipLayer := packet.Layer(layers.LayerTypeIPv4); ipLayer != nil {
		ip, _ := ipLayer.(*layers.IPv4)
		switch ip.Protocol {
		case layers.IPProtocolTCP:
			if tcpLayer := packet.Layer(layers.LayerTypeTCP); tcpLayer != nil {
				tcp, _ := tcpLayer.(*layers.TCP)
				handleTCPPacket(ip.SrcIP, ip.DstIP, tcp, packet)
			}
		case layers.IPProtocolUDP:
			if udpLayer := packet.Layer(layers.LayerTypeUDP); udpLayer != nil {
				udp, _ := udpLayer.(*layers.UDP)
				handleUDPPacket(ip.SrcIP, ip.DstIP, udp, packet)
			}
		case layers.IPProtocolICMPv4:
			if icmpLayer := packet.Layer(layers.LayerTypeICMPv4); icmpLayer != nil {
				icmp, _ := icmpLayer.(*layers.ICMPv4)
				handleICMPPacket(ip, icmp, packet)
			}
		}
	}

	// 解析IPv6
	if ip6Layer := packet.Layer(layers.LayerTypeIPv6); ip6Layer != nil {
		ip6, _ := ip6Layer.(*layers.IPv6)
		handleIPv6Packet(ip6, packet)
	}
}

// 打印全协议统计信息
func printStats() {
	flowMutex.RLock()
	defer flowMutex.RUnlock()

	// 统一定义列宽常量，便于维护和对齐
	const (
		connStrWidth    = 55
		protoWidth      = 8
		statusWidth     = 24
		durationWidth   = 12
		upKBWidth       = 20
		downKBWidth     = 20
		avgPktWidth     = 12
		lossStrWidth    = 12
		appNameWidth    = 20
		appCatWidth     = 15
		confidenceWidth = 6
	)

	fmt.Printf("\n===== 网络流监控统计 (%s) =====\n", time.Now().Format(time.RFC3339))
	fmt.Printf("总网络流数: %d (TCP/UDP/ICMP)\n", len(netFlows))
	fmt.Println("----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------")
	// 列头格式与数据格式使用相同的宽度定义
	fmt.Printf("%-*s %-*s %-*s %-*s %-*s %-*s %-*s %-*s %-*s %-*s %-*s \n",
		connStrWidth, "流信息",
		protoWidth, "协议",
		statusWidth, "状态",
		durationWidth, "连接时长",
		upKBWidth, "上行(KB)",
		downKBWidth, "下行(KB)",
		avgPktWidth, "平均包大小",
		lossStrWidth, "丢包/重传",
		appNameWidth, "应用名称",
		appCatWidth, "应用类别",
		confidenceWidth, "可信度")
	fmt.Println("----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------")

	flows := make([]*NetFlow, 0, len(netFlows))
	for _, flow := range netFlows {
		flows = append(flows, flow)
	}
	// 按最近活动时间排序
	sort.Slice(flows, func(i, j int) bool {
		return flows[i].LastSeen.After(flows[j].LastSeen)
	})

	displayCount := displayCountMax
	if len(flows) < displayCount {
		displayCount = len(flows)
	}

	for i := 0; i < displayCount; i++ {
		flow := flows[i]
		var connStr string
		if flow.Proto == "ICMP" || flow.Proto == "ICMPv6" {
			connStr = fmt.Sprintf("%s -> %s", flow.SrcIP, flow.DstIP)
		} else {
			connStr = fmt.Sprintf("%s:%d -> %s:%d", flow.SrcIP, flow.SrcPort, flow.DstIP, flow.DstPort)
		}
		// 截断过长的连接信息，避免撑宽列
		if len(connStr) > connStrWidth {
			connStr = connStr[:connStrWidth-3] + "..."
		}

		// 计算并格式化连接时长
		duration := flow.Duration()
		var durationStr string
		switch {
		case duration < time.Second:
			durationStr = fmt.Sprintf("%.0fms", duration.Seconds()*1000)
		case duration < time.Minute:
			durationStr = fmt.Sprintf("%.1fs", duration.Seconds())
		default:
			durationStr = fmt.Sprintf("%.1fm", duration.Minutes())
		}

		upKB := float64(flow.BytesSent) / 1024
		downKB := float64(flow.BytesReceived) / 1024
		lossStr := "-"

		switch flow.Proto {
		case "TCP":
			lossStr = fmt.Sprintf("丢:%d 重传:%d", flow.LostPackets, flow.RetransmittedPackets)
		case "UDP":
			lossStr = fmt.Sprintf("丢失:%d", flow.UDPDatagramLoss)
		case "ICMP", "ICMPv6":
			lossStr = fmt.Sprintf("类型:%d", flow.ICMPType)
		}
		// 截断过长的丢包信息
		if len(lossStr) > lossStrWidth {
			lossStr = lossStr[:lossStrWidth-3] + "..."
		}

		// 应用名称截断
		appName := flow.AppName
		if len(appName) > appNameWidth {
			appName = appName[:appNameWidth-3] + "..."
		}

		// 应用类别截断
		appCat := flow.AppCategory
		if len(appCat) > appCatWidth {
			appCat = appCat[:appCatWidth-3] + "..."
		}

		// 使用统一的宽度常量格式化输出，确保与列头对齐
		fmt.Printf("%-*s %-*s %-*s %-*s %-*.*f %-*.*f %-*.*f %-*s %-*s %-*s %-*.*f\n",
			connStrWidth, connStr,
			protoWidth, flow.Proto,
			statusWidth, flow.Status,
			durationWidth, durationStr,
			upKBWidth, 2, upKB,
			downKBWidth, 2, downKB,
			avgPktWidth, 1, flow.AvgPacketSize,
			lossStrWidth, lossStr,
			appNameWidth, appName,
			appCatWidth, appCat,
			confidenceWidth, 2, flow.Confidence)
	}

	fmt.Println("----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------")
}

// 清理过期流（按协议适配超时时间）
func cleanupOldFlows() {
	flowMutex.Lock()
	defer flowMutex.Unlock()

	now := time.Now()
	for key, flow := range netFlows {
		var expiration time.Duration
		switch flow.Proto {
		case "TCP":
			expiration = 5 * time.Minute // TCP连接超时较长
		case "UDP":
			expiration = udpTimeout
		case "ICMP", "ICMPv6":
			expiration = icmpTimeout
		default:
			expiration = 1 * time.Minute
		}

		if now.Sub(flow.LastSeen) > expiration {
			savePayloadsToFile(flow)
			delete(netFlows, key)
		}
	}
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func main() {
	handle, err := pcap.OpenLive(iface, snapshotLen, promiscuous, timeout)
	if err != nil {
		log.Fatalf("无法打开设备 %s: %v", iface, err)
	}
	defer handle.Close()

	// 设置过滤器：捕获TCP、UDP、ICMPv4、ICMPv6
	if err := handle.SetBPFFilter("tcp or udp or icmp or icmp6"); err != nil {
		log.Fatalf("无法设置过滤器: %v", err)
	}

	// 定时任务：打印统计、清理过期流、更新活跃流指标
	go func() {
		ticker := time.NewTicker(printInterval)
		defer ticker.Stop()
		for range ticker.C {
			printStats()
			cleanupOldFlows()
			updateActiveFlowsMetrics() // 新增：定期更新活跃流指标
		}
	}()

	// 处理退出信号
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
	packetChan := packetSource.Packets()

	fmt.Printf("开始监控网络流 on %s...\n", iface)
	fmt.Printf("监控协议: TCP, UDP, ICMPv4, ICMPv6\n")
	fmt.Printf("每 %v 打印一次统计信息\n", printInterval)
	fmt.Printf("Payload将存储在: %s\n", payloadStorageDir)
	fmt.Printf("Prometheus指标暴露在: http://%s/metrics\n", metricsAddr)
	fmt.Println("按Ctrl+C退出")

	// 主循环处理数据包
	for {
		select {
		case packet := <-packetChan:
			handlePacket(packet)
		case <-sigChan:
			fmt.Println("\n收到退出信号，正在保存最后的统计信息和payload...")
			// 保存所有活跃流的payload
			flowMutex.RLock()
			for _, flow := range netFlows {
				savePayloadsToFile(flow)
			}
			flowMutex.RUnlock()

			printStats()
			fmt.Println("监控已停止")
			return
		}
	}
}
